import os
import pymysql
from dbutils.pooled_db import PooledDB
from pymysql import cursors
from contextlib import contextmanager
from typing import Optional, Union, List, Dict, Any


class DBError(Exception):
    def __init__(self, message="数据库操作错误", sql=None, params=None):
        self.message = message
        self.sql = sql
        self.params = params
        super().__init__(self.message)

    def __str__(self):
        info = f"DBError: {self.message}"
        if self.sql:
            info += f"\nSQL: {self.sql}"
            if self.params:
                info += f"\nParams: {self.params}"
        return info

    pass


def _release_db_connection(conn=None):
    if conn:
        try:
            conn.close()
        except pymysql.Error as e:
            raise DBError(f"关闭数据库连接失败: {e}") from e


class DB:
    def __init__(self):
        self.pool = PooledDB(
            creator=pymysql,
            maxconnections=10,
            mincached=2,
            maxcached=5,
            blocking=True,
            setsession=[],
            ping=0,
            host='47.108.174.92',
            port=3306,
            user='bili',
            password=os.environ["MYSQL_USER_PASSWORD"],
            database='bilibili',
            charset='utf8mb4',
            cursorclass=cursors.DictCursor
        )

    @contextmanager
    def connection(self) -> pymysql.connections.Connection:
        conn = None
        try:
            conn = self._get_db_connection()
            yield conn
            conn.commit()
        except pymysql.Error as e:
            if conn:
                conn.rollback()
            raise DBError(f"数据库操作失败: {e}") from e
        finally:
            if conn:
                _release_db_connection(conn)

    @contextmanager
    def _get_cursor(self, conn=None):
        should_close_conn = False
        cursor = None

        try:
            if conn is None:
                conn = self.pool.connection()
                should_close_conn = True

            cursor = conn.cursor()
            yield cursor

            if should_close_conn:
                conn.commit()
        except pymysql.Error as e:
            if should_close_conn and conn:
                conn.rollback()
            raise DBError(f"数据库操作失败: {e}") from e
        finally:
            if cursor:
                cursor.close()
            if should_close_conn and conn:
                conn.close()

    def fetch_all(self, sql: str, params: Optional[Union[tuple, list, dict]] = None, conn=None) -> List[Dict[str, Any]]:
        with self._get_cursor() as cursor:
            cursor.execute(sql, params or ())
            return cursor.fetchall()

    def fetch_one(self, sql: str, params: Optional[Union[tuple, list, dict]] = None, conn=None) -> Optional[
        Dict[str, Any]]:
        with self._get_cursor(conn) as cursor:
            cursor.execute(sql, params or ())
            return cursor.fetchone()

    def execute(self, sql: str, params: Optional[Union[tuple, list, dict]] = None, conn=None) -> int:
        with self._get_cursor(conn) as cursor:
            cursor.execute(sql, params or ())
            return cursor.rowcount

    def _get_db_connection(self, conn=None):
        try:
            conn = self.pool.connection()
            return conn
        except pymysql.Error as e:
            raise DBError(f"获取数据库连接失败: {e}") from e


db = DB()
